package repository

import (
	"context"
	"go-pan/internal/model"
	"gorm.io/gorm"
)

type UploadChunkRepository interface {
	// 插入单个分片记录
	CreateChunk(ctx context.Context, chunk *model.UploadChunk) error

	UploadChunk(ctx context.Context, chunk *model.UploadChunk) error

	// 批量插入分片（可选）
	BatchCreateChunks(ctx context.Context, chunks []*model.UploadChunk) error

	// 查询某个上传任务的所有分片
	GetChunksByUploadID(ctx context.Context, uploadID int64) ([]*model.UploadChunk, error)

	// 判断某个分片是否已上传
	Exists(ctx context.Context, uploadID int64, chunkIndex int) (bool, error)

	// 统计某个任务上传的分片数量
	CountUploadedChunks(ctx context.Context, uploadID int64) (int, error)

	// 删除某个上传任务的所有分片（用于重传/回滚）
	DeleteChunksByUploadID(ctx context.Context, uploadID int64) error
}

type UploadChunkRepositoryImpl struct {
	db *gorm.DB
}

func NewUploadChunkRepository(db *gorm.DB) *UploadChunkRepositoryImpl {
	return &UploadChunkRepositoryImpl{db: db}
}

func (r *UploadChunkRepositoryImpl) CreateChunk(ctx context.Context, chunk *model.UploadChunk) error {
	return r.db.WithContext(ctx).Create(chunk).Error
}

func (r *UploadChunkRepositoryImpl) UploadChunk(ctx context.Context, chunk *model.UploadChunk) error {
	return r.db.WithContext(ctx).Save(chunk).Error
}

func (r *UploadChunkRepositoryImpl) BatchCreateChunks(ctx context.Context, chunks []*model.UploadChunk) error {
	return r.db.WithContext(ctx).Create(&chunks).Error
}

func (r *UploadChunkRepositoryImpl) GetChunksByUploadID(ctx context.Context, uploadID int64) ([]*model.UploadChunk, error) {
	var chunks []*model.UploadChunk
	if err := r.db.WithContext(ctx).Where("upload_id = ?", uploadID).Find(&chunks).Error; err != nil {
		return nil, err
	}
	return chunks, nil
}

func (r *UploadChunkRepositoryImpl) Exists(ctx context.Context, uploadID int64, chunkIndex int) (bool, error) {
	var count int64
	err := r.db.WithContext(ctx).Model(&model.UploadChunk{}).
		Where("upload_id = ? AND chunk_index = ?", uploadID, chunkIndex).
		Count(&count).Error
	if err != nil {
		return false, err
	}
	return count > 0, nil
}

func (r *UploadChunkRepositoryImpl) CountUploadedChunks(ctx context.Context, uploadID int64) (int, error) {
	var count int64
	err := r.db.WithContext(ctx).Model(&model.UploadChunk{}).
		Where("upload_id = ?", uploadID).Count(&count).Error
	if err != nil {
		return 0, err
	}
	return int(count), nil
}

func (r *UploadChunkRepositoryImpl) DeleteChunksByUploadID(ctx context.Context, uploadID int64) error {
	return r.db.WithContext(ctx).
		Where("upload_id = ?", uploadID).
		Delete(&model.UploadChunk{}).Error
}
